/* ====================================================================
 * Copyright (C) 2012 - Deniz Sezen
 * ====================================================================
 * File: dx_hook.cpp
 * Purpose: Contains the CFunctionHook class which is used to detour
 *	functions in Deus Ex: HR.
 * ====================================================================
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

/************************************************************************/
/* Includes                                                             */
/************************************************************************/

/* Operating system includes. */
#include <windows.h>

/* dxhr-sdk includes. */
#include "asm.h"
#include "dx_hook.h"

CFunctionHook::CFunctionHook(unsigned char* pTarget, void* pRedirect, int type)
{
	m_pTarget	= pTarget;
	m_pRedirect = pRedirect;
	m_iDetourType = 0;
}

CFunctionHook::~CFunctionHook( void )
{
	UndoPatch();

	delete m_pTrampoline;
	m_pTrampoline = NULL;
}

unsigned char* CFunctionHook::GetTrampoline( void ) 
{
	return m_pTrampoline;
}

DWORD CFunctionHook::MakeMemWritable(void* target, unsigned int size, DWORD prot)
{
	// Unprotect the function bytes.
	DWORD old_prot;
	VirtualProtect(target, size, prot, &old_prot);
	return old_prot;
}

bool CFunctionHook::DoPatch_RET( void )
{
	// push functocall
	// 68 XX XX XX XX
	
	// ret
	// C3

	// Need 6 bytes.
	m_nSavedBytes = copy_bytes(m_pTarget, NULL, 6);

	// Allocate space for the trampoline.
	// Need 6 + 5 more to jump back.
	m_pTrampoline = new unsigned char[m_nSavedBytes + JMP_SIZE];
	MakeMemWritable(m_pTrampoline, m_nSavedBytes + JMP_SIZE, PAGE_EXECUTE_READWRITE);
	memset(m_pTrampoline, 0x90, m_nSavedBytes + JMP_SIZE);

	// Copy function bytes.
	copy_bytes(m_pTarget, m_pTrampoline, m_nSavedBytes);

	// Inject a jump at the end to the origfunc + m_nSavedBytes.
	unsigned char* src = m_pTrampoline + m_nSavedBytes;
	unsigned char* dest = m_pTarget + m_nSavedBytes;
	MakeMemWritable(src, JMP_SIZE, PAGE_EXECUTE_READWRITE);
	inject_jmp(src, dest);

	// Give us permission to write in the opcodes.
	DWORD old_prot = MakeMemWritable(m_pTarget, m_nSavedBytes, PAGE_EXECUTE_READWRITE);

	// Write the opcodes.
	m_pTarget[0] = '\x68';
	*(long *)&m_pTarget[1] = (long)m_pRedirect;
	m_pTarget[5] = '\xC3';

	// Reset memory page permissions.
	MakeMemWritable(m_pTarget, m_nSavedBytes, old_prot);

	// Flush processor instruction cache to make sure we have the changed bytes in
	// memory.
	FlushInstructionCache( GetCurrentProcess(), m_pTarget, m_nSavedBytes);

	return true;
}

bool CFunctionHook::DoPatch_JMP( void )
{
	if( !m_pTarget || !m_pRedirect ) {
		return false;
	}

	// Figure out the number of bytes we need to save away.
	m_nSavedBytes = copy_bytes(m_pTarget, NULL, JMP_SIZE);

	// Allocate space for the trampoline.
	m_pTrampoline = new unsigned char[m_nSavedBytes + JMP_SIZE];
	memset(m_pTrampoline, 0x90, m_nSavedBytes + JMP_SIZE);

	// Make sure we can write to it and copy the bytes.
	MakeMemWritable(m_pTrampoline, m_nSavedBytes + JMP_SIZE, PAGE_EXECUTE_READWRITE);
	copy_bytes(m_pTarget, m_pTrampoline, JMP_SIZE);

	// Place a jump back to the target+5 into the trampoline.
	unsigned char* src = m_pTrampoline + m_nSavedBytes;
	unsigned char* dest = m_pTarget + m_nSavedBytes;
	MakeMemWritable(src, JMP_SIZE, PAGE_EXECUTE_READWRITE);
	inject_jmp(src, dest);

	// Make sure the target function memory is writeable and
	// inject the jump to that as well.
	DWORD old_prot = MakeMemWritable(m_pTarget, m_nSavedBytes, PAGE_EXECUTE_READWRITE);
	inject_jmp((void *)m_pTarget, m_pRedirect);
	MakeMemWritable(m_pTarget, m_nSavedBytes, old_prot);

	return true;
}

void CFunctionHook::UndoPatch( void )
{
	// Make sure we can write to the target.
	DWORD old_prot = MakeMemWritable(m_pTarget, m_nSavedBytes, PAGE_EXECUTE_READWRITE);

	// Copy back the saved bytes.
	copy_bytes(m_pTrampoline, m_pTarget, m_nSavedBytes);

	// Reset page permissions.
	MakeMemWritable(m_pTarget, m_nSavedBytes, old_prot);

	// Flush the instruction cache.
	FlushInstructionCache(GetCurrentProcess(), m_pTarget, m_nSavedBytes);
}

